#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <memory.h>
#include <time.h>
#include <float.h>
#include <stdbool.h>

#include "ops.h"
#include "ops_eig_sol_gcg.h"
#include "app_slepc.h"
#if 1
int TestEigenSolverGCG   (void *A, void *B, int flag, int argc, char *argv[], struct OPS_ *ops, double shift);
int TestAppSLEPC (int argc, char *argv[]);

static char help[] = "Test App of SLEPC.\n";
int TestEigenSolverGCG(void *A, void *B, int flag, int argc, char *argv[], struct OPS_ *ops, double shift) 
{

	int nevConv  = 10; //number of required eigenpairs 
	double gapMin = 1e-5, abs_tol, rel_tol;
	int nevGiven = 0, block_size = nevConv/2, nevMax = 2*nevConv, multiMax = 1, nevInit = nevMax;
	abs_tol = 1e-7; //abs tolerance
	rel_tol = 1e-9; //rel tolerance
	//to choose the method for linear solver
	// char compW_method[8] = "bcg"; //mumps,superludist,petsc_pcg,pardiso

	if (nevConv > 300)
	{
		block_size = nevConv / 10;
		nevInit = 5*block_size; 
		nevMax =  nevConv + nevInit;
	}

	ops->GetOptionFromCommandLine("-nevConv"  ,'i',&nevConv   ,argc,argv,ops);
	ops->GetOptionFromCommandLine("-nevMax"   ,'i',&nevMax    ,argc,argv,ops);
	ops->GetOptionFromCommandLine("-blockSize",'i',&block_size,argc,argv,ops);
	ops->GetOptionFromCommandLine("-nevInit"  ,'i',&nevInit   ,argc,argv,ops);
	nevConv= atof(argv[5]);
	block_size=atof(argv[6]);
	nevInit=atof(argv[7]);
	nevMax=atof(argv[8]);

	nevInit = nevInit<nevMax?nevInit:nevMax;
	// nevGiven = nevInit;
	int max_iter_gcg = 400; 
	double tol_gcg[2] = {1e-7,1e-9};
	rel_tol=atof(argv[9]);
	abs_tol=1.0;
	tol_gcg[0] = abs_tol;
	tol_gcg[1] = rel_tol;

	
	double *eval; void **evec;
	eval = malloc(nevMax*sizeof(double));
	memset(eval,0,nevMax*sizeof(double));
	ops->MultiVecCreateByMat(&evec,nevMax,A,ops);
	ops->MultiVecSetRandomValue(evec,0,nevMax,ops);
	void **gcg_mv_ws[4]; double *dbl_ws; int *int_ws;
	ops->MultiVecCreateByMat(&gcg_mv_ws[0],nevMax+2*block_size,A,ops);				
	ops->MultiVecSetRandomValue(gcg_mv_ws[0],0,nevMax+2*block_size,ops);
	ops->MultiVecCreateByMat(&gcg_mv_ws[1],block_size,A,ops);				
	ops->MultiVecSetRandomValue(gcg_mv_ws[1],0,block_size,ops);
	ops->MultiVecCreateByMat(&gcg_mv_ws[2],block_size,A,ops);				
	ops->MultiVecSetRandomValue(gcg_mv_ws[2],0,block_size,ops);
	ops->MultiVecCreateByMat(&gcg_mv_ws[3],block_size,A,ops);				
	ops->MultiVecSetRandomValue(gcg_mv_ws[3],0,block_size,ops);
	// BV evecBV, VBV;
	// LoadBV(PETSC_COMM_WORLD, "/share/home/guozhaotong/ztguo/GCGE0/build/X0.bin", &evecBV);
	// LoadBV(PETSC_COMM_WORLD, "/share/home/guozhaotong/ztguo/GCGE0/build/V.bin", &VBV);

	// Mat evecMat,evecBVMat;
	// BVSetActiveColumns((BV)evec, 0, nevMax);
	// BVSetActiveColumns(evecBV, 0, nevMax);
	// BVGetMat((BV)evec,&evecMat);
	// BVGetMat(evecBV, &evecBVMat);
	// MatCopy(evecBVMat, evecMat,  DIFFERENT_NONZERO_PATTERN);
	// BVRestoreMat((BV)evec,&evecMat);
	// BVRestoreMat(evecBV, &evecBVMat);

	// BVSetActiveColumns((BV)gcg_mv_ws[0], 0, nevMax+2*block_size);
	// BVSetActiveColumns(VBV, 0, nevMax+2*block_size);
	// BVGetMat((BV)gcg_mv_ws[0],&evecMat);
	// BVGetMat(VBV, &evecBVMat);
	// MatCopy(evecBVMat, evecMat,  DIFFERENT_NONZERO_PATTERN);
	// BVRestoreMat((BV)gcg_mv_ws[0],&evecMat);
	// BVRestoreMat(VBV, &evecBVMat);

	int sizeV = nevInit + 2*block_size;
	int length_dbl_ws = 2*sizeV*sizeV+10*sizeV
		+(nevMax+2*block_size)+(nevMax)*block_size;
	// ops->Printf ( "length_dbl_ws = %d\n", length_dbl_ws );
	int length_int_ws = 6*sizeV+2*(block_size+3);
	// ops->Printf ( "length_int_ws = %d\n", length_int_ws );
	dbl_ws = malloc(length_dbl_ws*sizeof(double));
	memset(dbl_ws,0,length_dbl_ws*sizeof(double));
	int_ws = malloc(length_int_ws*sizeof(int));
	memset(int_ws,0,length_int_ws*sizeof(int));

	srand(0);
	double time_start, time_interval;
	time_start = ops->GetWtime();
		
	ops->Printf("===============================================\n");
	ops->Printf("GCG Eigen Solver\n");
	EigenSolverSetup_GCG(multiMax,gapMin,nevInit,nevMax,block_size,
		tol_gcg,max_iter_gcg,flag,gcg_mv_ws,dbl_ws,int_ws,ops);
	
	int    check_conv_max_num    = 50   ;
		
	char   initX_orth_method[8]  = "mgs"; 
	int    initX_orth_block_size = 80   ; 
	int    initX_orth_max_reorth = 2    ; double initX_orth_zero_tol    = 2*DBL_EPSILON;//1e-12
	
	char   compP_orth_method[8]  = "mgs"; 
	int    compP_orth_block_size = -1   ; 
	int    compP_orth_max_reorth = 2    ; double compP_orth_zero_tol    = 2*DBL_EPSILON;//1e-12
	
	char   compW_orth_method[8]  = "mgs";
	int    compW_orth_block_size = 80   ; 	
	int    compW_orth_max_reorth = 2    ;  double compW_orth_zero_tol   = 2*DBL_EPSILON;//1e-12
	int    compW_bpcg_max_iter   = 500   ;  double compW_bpcg_rate       = 1e-2; 
	double compW_bpcg_tol        = 1e-14;  char   compW_bpcg_tol_type[8] = "abs";
	
	int    compRR_min_num        = -1   ;  double compRR_min_gap        = gapMin;
	double compRR_tol            = 2*DBL_EPSILON;
	bool   compW_pc				 = false;	
	int	   auto_shift			 =1;
	compW_bpcg_max_iter=atof(argv[10]);
	compW_pc		   =atof(argv[11]);
	char compW_method[20] = "bpcg"; //mumps,superludist,pardiso 有高优先级
	auto_shift      = atof(argv[12])   ;
	
	EigenSolverSetParameters_GCG(
			check_conv_max_num   ,
			initX_orth_method    , initX_orth_block_size, 
			initX_orth_max_reorth, initX_orth_zero_tol  ,
			compP_orth_method    , compP_orth_block_size, 
			compP_orth_max_reorth, compP_orth_zero_tol  ,
			compW_orth_method    , compW_orth_block_size, 
			compW_orth_max_reorth, compW_orth_zero_tol  ,
			compW_bpcg_max_iter  , compW_bpcg_rate      , 
			compW_bpcg_tol       , compW_bpcg_tol_type  ,
			auto_shift, /* auto_shift */
			compRR_min_num       , compRR_min_gap       ,
			compRR_tol           , compW_method         , 
			compW_pc,
			ops);		

	EigenSolverSetParametersFromCommandLine_GCG(argc,argv,ops);
	ops->Printf("nevGiven = %d, nevConv = %d, nevMax = %d, block_size = %d, nevInit = %d\n",
			nevGiven,nevConv,nevMax,block_size,nevInit);

	ops->EigenSolver(A,B,eval,evec,nevGiven,&nevConv,ops);
	ops->Printf("numIter = %d, nevConv = %d\n",
			((GCGSolver*)ops->eigen_solver_workspace)->numIter, nevConv);
	ops->Printf("++++++++++++++++++++++++++++++++++++++++++++++\n");

	time_interval = ops->GetWtime() - time_start;
	ops->Printf("Time is %.8f\n", time_interval);

	ops->MultiVecDestroy(&gcg_mv_ws[0],nevMax+2*block_size,ops);
	ops->MultiVecDestroy(&gcg_mv_ws[1],block_size,ops);
	ops->MultiVecDestroy(&gcg_mv_ws[2],block_size,ops);
	ops->MultiVecDestroy(&gcg_mv_ws[3],block_size,ops);

	free(dbl_ws); free(int_ws);

#if 0
	ops->Printf("eigenvalues\n");
	int idx;
	for (idx = 0; idx < nevConv; ++idx) {
		eval[idx] -= shift;
		ops->Printf("%.10f\n",eval[idx]);
	}
#endif
	ops->MultiVecDestroy(&(evec),nevMax,ops);
	free(eval);
	return 0;
}

int TestAppSLEPC(int argc, char *argv[]) 
{
	
	SlepcInitialize(&argc,&argv,(char*)0,help);
	PetscMPIInt   rank, size;
  	MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
  	MPI_Comm_size(PETSC_COMM_WORLD, &size);	
	
	OPS *slepc_ops = NULL;
	OPS_Create (&slepc_ops);
	OPS_SLEPC_Set (slepc_ops);
	OPS_Setup (slepc_ops);
	slepc_ops->Printf("%s", help);
	slepc_ops->Printf("================================================\n");
	slepc_ops->Printf("===========(0)Load Mat K and M===========\n");
	
	void *matA, *matB; OPS *ops;
	double t0, t1;

	t0 = MPI_Wtime();

   	Mat      slepc_matA, slepc_matB;
   	PetscBool flg;
	PetscInt nrows, ncols;
#if 0
	PetscViewer    viewer;
	PetscViewerBinaryOpen(PETSC_COMM_WORLD,"bigK.petsc.bin",FILE_MODE_READ,&viewer); 
	MatCreate(PETSC_COMM_WORLD,&slepc_matA); 
	MatSetFromOptions(slepc_matA);
	MatLoad(slepc_matA,viewer); 
	PetscViewerDestroy(&viewer);
	MatGetSize(slepc_matA,&nrows,&ncols);
	slepc_ops->Printf("matrix A %d, %d\n",nrows,ncols);
	
	PetscViewer    viewer1;
	PetscViewerBinaryOpen(PETSC_COMM_WORLD,"bigM.petsc.bin",FILE_MODE_READ,&viewer1); 
	MatCreate(PETSC_COMM_WORLD,&slepc_matB); 
	MatSetFromOptions(slepc_matB);
	MatLoad(slepc_matB,viewer1); 
	PetscViewerDestroy(&viewer1);
	MatGetSize(slepc_matB,&nrows,&ncols);
	slepc_ops->Printf("matrix B %d, %d\n",nrows,ncols);
#else
	char file_matK[PETSC_MAX_PATH_LEN];
	strcpy(file_matK, argv[1]);
	slepc_ops->Printf("matK file : %s\n",file_matK);

	PetscViewer fd1;
	PetscViewerBinaryOpen(PETSC_COMM_WORLD, file_matK, FILE_MODE_READ, &fd1);
	MatCreate(PETSC_COMM_WORLD, &slepc_matA); 
	MatSetFromOptions(slepc_matA);
	MatLoad(slepc_matA,fd1); 
	PetscViewerDestroy(&fd1);
	MatGetSize(slepc_matA,&nrows,&ncols);
	slepc_ops->Printf("size of matrix K : %d, %d\n",nrows,ncols);
	
	char file_matM[PETSC_MAX_PATH_LEN];
	strcpy(file_matM, argv[2]);
	slepc_ops->Printf("matM file : %s\n",file_matM);

	PetscViewer fd2;
	PetscViewerBinaryOpen(PETSC_COMM_WORLD, file_matM, FILE_MODE_READ, &fd2);
	MatCreate(PETSC_COMM_WORLD, &slepc_matB); 
	MatSetFromOptions(slepc_matB);
	MatLoad(slepc_matB,fd2); 
	PetscViewerDestroy(&fd2);
	MatGetSize(slepc_matB,&nrows,&ncols);
	slepc_ops->Printf("size of matrix M : %d, %d\n",nrows,ncols);

#endif
	t1 = MPI_Wtime();
	slepc_ops->Printf("Time of Loading K and M   : %.2f s\n", t1 - t0);
	slepc_ops->Printf("===========(1)pre-treat Mat K and M===========\n");

	t0 = MPI_Wtime();

    Vec x, min, max;
    PetscInt idx;
    PetscReal max_num, min_num;
    MatCreateVecs(slepc_matA, &x, NULL);
    VecDuplicate(x, &max);
    VecDuplicate(x, &min);
    MatGetRowMaxAbs(slepc_matA, max, NULL);
    MatGetRowMinAbs(slepc_matA, min, NULL);
    VecMax(max, &idx, &max_num);
    VecMin(min, &idx, &min_num);
    slepc_ops->Printf("observe the elements of the given matrix\n");
    slepc_ops->Printf("max = %4.e, min = %4.e\n", max_num, min_num);
    VecDestroy(&max);
    VecDestroy(&min);
    VecDestroy(&x);

	/* 预处理 */
	int pre_treat_method = 0;
	pre_treat_method=atof(argv[3]);
#if 1
	if(pre_treat_method)
	{
		
		Vec diagK;
		Mat D_bigK;
		MatCreate(PETSC_COMM_WORLD, &D_bigK);
		MatSetSizes(D_bigK, PETSC_DECIDE, PETSC_DECIDE, nrows, ncols);
		MatSetUp(D_bigK);

		VecCreate(PETSC_COMM_WORLD, &diagK);
		VecSetSizes(diagK, PETSC_DECIDE, nrows);
		VecSetFromOptions(diagK);
		if (pre_treat_method==10)
		{
			slepc_ops->Printf("对角归一预条件\n");
			//construct D = diag(1./sqrt(diag(K))
			/* 对角归一预条件 */
			MatGetDiagonal(slepc_matA, diagK);	
		}
		else if(pre_treat_method==1)/* L1范数 */
		{
			slepc_ops->Printf("L1-预条件\n");
			PetscInt rstart, rend, ncols_row; /* 局部行号范围，每个行的非零元 */
			const PetscInt *cols; /* 每个元素列号 */
			const PetscScalar *vals; /* 每行的元素值 */
			PetscScalar norm_val;
			MatGetOwnershipRange(slepc_matA, &rstart, &rend);
			for(int ii=rstart;ii<rend;ii++)/* 逐行计算 */
			{
				MatGetRow(slepc_matA, ii, &ncols_row, &cols, &vals);
				norm_val = 0.0;
				for (PetscInt jj = 0; jj < ncols_row; jj++) {
        			norm_val += PetscAbsScalar(vals[jj]); /* L1范数 */
    			}
				VecSetValue(diagK, ii, norm_val, INSERT_VALUES);
				MatRestoreRow(slepc_matA, ii, &ncols_row, &cols, &vals);
			}

		}
		else if(pre_treat_method==2) /* L2范数 */
		{
			slepc_ops->Printf("L2-预条件\n");
			PetscInt rstart, rend, ncols_row; /* 局部行号范围，每个行的非零元 */
			const PetscInt *cols; /* 每个元素列号 */
			const PetscScalar *vals; /* 每行的元素值 */
			PetscScalar norm_val;
			MatGetOwnershipRange(slepc_matA, &rstart, &rend);
			for(int ii=rstart;ii<rend;ii++)/* 逐行计算 */
			{
				MatGetRow(slepc_matA, ii, &ncols_row, &cols, &vals);
				norm_val = 0.0;
				for (PetscInt jj = 0; jj < ncols_row; jj++) {
        			norm_val += PetscAbsScalar(vals[jj]) * PetscAbsScalar(vals[jj]); /* L2范数 */
    			}
				VecSetValue(diagK, ii, norm_val, INSERT_VALUES);
				MatRestoreRow(slepc_matA, ii, &ncols_row, &cols, &vals);
			}
		}
		else /* L∞ 范数 */
		{
			slepc_ops->Printf("L无穷-预条件\n");
			PetscInt rstart, rend, ncols_row; /* 局部行号范围，每个行的非零元 */
			const PetscInt *cols; /* 每个元素列号 */
			const PetscScalar *vals; /* 每行的元素值 */
			PetscScalar norm_val;
			PetscScalar absval;
			MatGetOwnershipRange(slepc_matA, &rstart, &rend);
			for(int ii=rstart;ii<rend;ii++)/* 逐行计算 */
			{
				MatGetRow(slepc_matA, ii, &ncols_row, &cols, &vals);
				norm_val = 0.0;
				for (PetscInt jj = 0; jj < ncols_row; jj++) 
				{
					absval = PetscAbsScalar(vals[jj]);
					if (absval > norm_val) norm_val = absval; /* L∞ 范数 */
    			}
				VecSetValue(diagK, ii, norm_val, INSERT_VALUES);
				MatRestoreRow(slepc_matA, ii, &ncols_row, &cols, &vals);
			}
		}

		VecSqrtAbs(diagK);
		VecPow(diagK, -1);
		MatDiagonalSet(D_bigK, diagK, INSERT_VALUES);
		MatAssemblyBegin(D_bigK, MAT_FINAL_ASSEMBLY);
    	MatAssemblyEnd(D_bigK, MAT_FINAL_ASSEMBLY);


		MatMatMatMult(D_bigK, slepc_matA, D_bigK,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&slepc_matA);
		MatMatMatMult(D_bigK, slepc_matB, D_bigK,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&slepc_matB);
		VecDestroy(&diagK);
		MatDestroy(&D_bigK);
		
	}
	else
	{
		PetscPrintf(PETSC_COMM_WORLD,"未使用预处理\n");
		// MatAssemblyBegin(slepc_matB, MAT_FINAL_ASSEMBLY);
		// MatAssemblyEnd(slepc_matB, MAT_FINAL_ASSEMBLY);
	}

#endif 

	

#if 1
    double sig = 0;
	sig=atof(argv[3]);

	MatAXPY(slepc_matA,sig,slepc_matB,DIFFERENT_NONZERO_PATTERN);

#endif
	MatAssemblyBegin(slepc_matA, MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(slepc_matA, MAT_FINAL_ASSEMBLY);

	t1 = MPI_Wtime();
	slepc_ops->Printf("Time of Pre-treat K and M : %.2f s\n", t1 - t0);
	slepc_ops->Printf("Shift                     : %f\n",sig);
	slepc_ops->Printf("================================================\n");

	ops = slepc_ops; 
	matA = (void*)(slepc_matA); matB = (void*)(slepc_matB);
	int flag = 0;
	TestEigenSolverGCG(matA,matB,flag,argc,argv,ops,sig);

		
	/* 销毁petsc矩阵 */
   	MatDestroy(&slepc_matA);
   	MatDestroy(&slepc_matB);
	OPS_Destroy (&slepc_ops);
	SlepcFinalize();
	
	return 0;
}
#endif
int main(int argc, char *argv[]) 
{
	TestAppSLEPC(argc, argv);
   	return 0;
}
